package client_core

import (
	"context"
	"encoding/json"
	"github.com/chicken-team-outside/chicken_transmission/client_core/service"
	"github.com/chicken-team-outside/chicken_transmission/utils"
	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
	"go.uber.org/zap"
	"io"
	"sync"
)

type RpcConnectionStatus int8

var upgrader *websocket.Upgrader

var clientMap map[uint64]*ISocketConnection

var clientMapLock sync.RWMutex

var clientId uint64

var clientIdLock sync.Mutex

type EventRequest struct {
	Id      int64           `json:"id"`
	Event   string          `json:"event"`
	Content json.RawMessage `json:"content"`
}

type EventReply struct {
	Id     int64         `json:"id"`
	Event  string        `json:"event,omitempty"`
	Err    string        `json:"err,omitempty"`
	Result []interface{} `json:"result,omitempty"`
}

type ISocketConnection struct {
	conn       *websocket.Conn
	ctx        context.Context
	cancel     context.CancelCauseFunc
	msgHandler func(*ISocketConnection, []byte) error
	sendQueue  *utils.BlockingQueue
	closeHold  sync.Once
}

func NewISocketConnection(ctx context.Context, conn *websocket.Conn, onMsg func(*ISocketConnection, []byte) error, onClose func()) (r *ISocketConnection) {
	r = &ISocketConnection{
		conn:       conn,
		msgHandler: onMsg,
	}
	conn.SetCloseHandler(func(code int, text string) error {
		r.cancel(websocket.ErrCloseSent)
		onClose()
		return nil
	})
	r.ctx, r.cancel = context.WithCancelCause(ctx)
	r.sendQueue = utils.NewBlockingQueue(r.ctx)
	return
}

func (c *ISocketConnection) runReadMsg() (err error) {
	defer c.cancel(err)
	var r io.Reader
	for c.ctx.Err() == nil {
		_, r, err = c.conn.NextReader()
		var data []byte
		if err == nil {
			data, err = io.ReadAll(r)
		}
		if err == nil {
			err = c.msgHandler(c, data)
		}
		if err != nil {
			return
		}
	}
	return
}

func (c *ISocketConnection) runWriteMsg() (err error) {
	defer c.cancel(err)
	var data interface{}
	var writer io.WriteCloser
	tmp := make(chan interface{}, 1)
	for {
		data, err = c.sendQueue.Pop(tmp)
		if err != nil {
			return
		}
		writer, err = c.conn.NextWriter(websocket.TextMessage)
		if err == nil {
			_, err = writer.Write(data.([]byte))
			writer.Close()
		}
		if err != nil {
			return
		}
	}
}

func (c *ISocketConnection) Close() {
	c.closeHold.Do(func() {
		c.cancel(context.Canceled)
		c.conn.Close()
	})
}

func (c *ISocketConnection) SendMsg(data []byte) error {
	return c.sendQueue.Push(data)
}

func (c *ISocketConnection) SendEventReply(id int64, err error, args ...interface{}) {
	reply := &EventReply{
		Id:     id,
		Result: args,
	}
	if err != nil {
		reply.Err = err.Error()
	}
	marshal, err := json.Marshal(reply)
	if err != nil {
		zap.L().Error("json marshal error", zap.Error(err))
		return
	}
	c.SendMsg(marshal)
}

func initNotifyMgr(ctx context.Context, router *gin.RouterGroup) {
	router.Any("", func(c *gin.Context) {
		websocketConnectRequestHandler(ctx, c)
	})
	clientMap = make(map[uint64]*ISocketConnection)
	upgrader = &websocket.Upgrader{
		ReadBufferSize:  1024,
		WriteBufferSize: 1024,
	}
}

func getClientId() uint64 {
	clientIdLock.Lock()
	defer clientIdLock.Unlock()
	clientId++
	return clientId
}

func websocketConnectRequestHandler(rootCtx context.Context, ctx *gin.Context) {
	id := getClientId()
	conn, err := upgrader.Upgrade(ctx.Writer, ctx.Request, nil)
	if err == nil {
		zap.L().Debug("client connect", zap.String("RemoteAddr", ctx.Request.RemoteAddr))
		client := NewISocketConnection(rootCtx, conn, clientMsgHandler, func() {
			zap.L().Debug("client disconnect", zap.String("RemoteAddr", ctx.Request.RemoteAddr))
			clientMapLock.Lock()
			delete(clientMap, id)
			clientMapLock.Unlock()
		})
		clientMapLock.Lock()
		clientMap[id] = client
		clientMapLock.Unlock()
		go client.runReadMsg()
		go client.runWriteMsg()
	}
}

func clientMsgHandler(conn *ISocketConnection, data []byte) error {
	event := new(EventRequest)
	err := json.Unmarshal(data, event)
	var result interface{}
	if err == nil {
		switch event.Event {
		case "load-server-status":
			result, err = loadServerStatus(conn.ctx, event)
		case "ping":
			result = nil
		case "download-task-progress":
			result = downloadTaskRunnerMgr.GetRunningTaskStatus()
		case "upload-task-progress":
			result = uploadTaskRunnerMgr.GetRunningTaskStatus()
		default:
			zap.L().Debug("unknown event", zap.String("event", event.Event))
			return nil
		}
		if result == nil {
			conn.SendEventReply(event.Id, err)
		} else {
			conn.SendEventReply(event.Id, err, result)
		}
	}
	return err
}

func loadServerStatus(ctx context.Context, msg *EventRequest) (interface{}, error) {
	tasks, err := service.FindAllTask()
	if err != nil {
		return nil, err
	}
	return map[string]interface{}{
		"status": getCurrentStatus(ctx),
		"tasks":  tasks,
	}, nil
}

func sendEventData(event string, args ...interface{}) {
	clientMapLock.RLock()
	defer clientMapLock.RUnlock()
	reply := &EventReply{
		Event:  event,
		Result: args,
	}
	marshal, err := json.Marshal(reply)
	if err == nil {
		for _, client := range clientMap {
			go client.SendMsg(marshal)
		}
	}
}

func sendRpcStatusChange(status ServiceStatus) {
	sendEventData("rpc-status-change", status)
}
